# 导入必要的包
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.keras.models import Sequential, load_model, save_model
from tensorflow.keras.layers import LSTM, Dense
from sklearn.preprocessing import MinMaxScaler
import pickle
import os
from datetime import datetime
import matplotlib.dates as mdates

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


class PressureForecaster:
    """常州气压预测系统"""

    def __init__(self, data_file='常州.csv', time_step=10, epochs=20, batch_size=32):
        """
        初始化气压预测系统

        参数:
            data_file: 数据文件路径
            time_step: 时间步长，用于LSTM模型
            epochs: 训练轮次
            batch_size: 批量大小
        """
        self.data_file = data_file
        self.time_step = time_step
        self.epochs = epochs
        self.batch_size = batch_size

        # 提取文件名作为基础名称（不含路径和扩展名）
        self.base_filename = os.path.splitext(os.path.basename(data_file))[0]

        # 模型文件路径
        self.model_path = f'{self.base_filename}_lstm_model.h5'
        self.scaler_path = f'{self.base_filename}_scaler.pkl'
        self.info_path = f'{self.base_filename}_model_info.pkl'

        # 初始化变量
        self.df = None
        self.model = None
        self.scaler = None
        self.model_info = None

        # 列名映射
        self.column_mapping = {
            'temperture': 'temperature',  # 保持原拼写以匹配输入数据
            '气压': 'pressure',
            '降雨': 'rainfall',
            '风速': 'wind_speed',
            'date': 'date'  # 日期列
        }

        print(f"初始化常州气压预测系统")
        print(f"数据文件: {data_file}")
        print(f"模型文件将保存为: {self.model_path}")

    def load_data(self):
        """加载常州气象数据"""
        try:
            # 读取CSV文件
            df = pd.read_csv(self.data_file)
            # 修正列名（去除空格）
            df.columns = [col.strip() for col in df.columns]

            print(f"原始列名: {list(df.columns)}")

            # 确保日期列存在
            if 'date' not in df.columns:
                raise ValueError("数据文件缺少'date'列")

            # 将日期列转换为datetime类型并设为索引
            df['date'] = pd.to_datetime(df['date'])
            df.set_index('date', inplace=True)
            # 按时间排序
            df.sort_index(inplace=True)

            # 重命名列以匹配预期名称
            for orig_col, new_col in self.column_mapping.items():
                if orig_col in df.columns:
                    df[new_col] = df[orig_col]
                else:
                    print(f"警告: 找不到'{orig_col}'列")

            # 确保气压数据可用
            if 'pressure' not in df.columns and '气压' in df.columns:
                df['pressure'] = df['气压']
            elif 'pressure' not in df.columns:
                print("警告: 找不到气压数据，创建随机数据")
                df['pressure'] = np.random.normal(1013, 10, len(df))

            self.df = df
            print(f"成功加载常州数据，共 {len(df)} 条记录")
        except Exception as e:
            print(f"加载常州数据出错: {e}")
            print("使用生成的示例数据代替")
            self.df = self._generate_example_data()

        return self.df

    def _generate_example_data(self):
        """生成示例气压数据（内部方法）"""
        date_rng = pd.date_range(start='2023-01-01', end='2023-12-31', freq='D')
        pressure = np.random.normal(loc=1013, scale=10, size=len(date_rng))

        df = pd.DataFrame({'date': date_rng, 'pressure': pressure})
        df.set_index('date', inplace=True)
        print("已生成示例数据")
        return df

    def _create_dataset(self, dataset):
        """创建时间序列数据集（内部方法）"""
        X, Y = [], []
        for i in range(len(dataset) - self.time_step - 1):
            a = dataset[i:(i + self.time_step), 0]
            X.append(a)
            Y.append(dataset[i + self.time_step, 0])
        return np.array(X), np.array(Y)

    def train_model(self):
        """训练LSTM模型并保存"""
        print("开始训练气压预测模型...")

        # 确保数据已加载
        if self.df is None:
            self.load_data()

        # 确保数据中没有缺失值
        self.df = self.df.dropna(subset=['pressure'])

        # 数据预处理
        self.scaler = MinMaxScaler(feature_range=(0, 1))
        scaled_data = self.scaler.fit_transform(self.df[['pressure']])

        # 创建训练数据
        X, Y = self._create_dataset(scaled_data)
        X = X.reshape(X.shape[0], X.shape[1], 1)

        # 构建LSTM模型
        self.model = Sequential()
        self.model.add(LSTM(50, return_sequences=True, input_shape=(self.time_step, 1)))
        self.model.add(LSTM(50))
        self.model.add(Dense(1))
        self.model.compile(loss='mean_squared_error', optimizer='adam')

        # 训练模型
        history = self.model.fit(X, Y, epochs=self.epochs,
                                 batch_size=self.batch_size, verbose=1)
        print(f"LSTM模型训练完成，共训练 {self.epochs} 轮")

        # 保存模型信息
        self.model_info = {
            'time_step': self.time_step,
            'last_date': self.df.index[-1],
            'trained_date': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }

        # 保存模型
        self.save_model()

        return history

    def save_model(self):
        """保存训练好的模型和相关信息"""
        # 保存LSTM模型
        save_model(self.model, self.model_path)
        print(f"LSTM模型已保存到 {self.model_path}")

        # 保存scaler
        with open(self.scaler_path, 'wb') as f:
            pickle.dump(self.scaler, f)
        print(f"数据标准化器已保存到 {self.scaler_path}")

        # 保存模型信息
        with open(self.info_path, 'wb') as f:
            pickle.dump(self.model_info, f)
        print(f"模型信息已保存到 {self.info_path}")

    def load_model(self):
        """加载保存的模型"""
        try:
            # 加载LSTM模型
            self.model = load_model(self.model_path)

            # 加载scaler
            with open(self.scaler_path, 'rb') as f:
                self.scaler = pickle.load(f)

            # 加载模型信息
            with open(self.info_path, 'rb') as f:
                self.model_info = pickle.load(f)

            print("模型加载成功")
            print(f"模型训练日期: {self.model_info.get('trained_date', '未知')}")
            return True
        except Exception as e:
            print(f"加载模型时出错: {e}")
            return False

    def predict(self, days=7):
        """预测未来气压"""
        # 检查是否需要加载数据
        if self.df is None:
            self.load_data()

        # 检查是否需要加载或训练模型
        if self.model is None or self.scaler is None:
            if os.path.exists(self.model_path) and self.load_model():
                print("使用已保存的模型进行预测")
            else:
                print("未找到模型或加载失败，训练新模型...")
                self.train_model()

        # 提取时间步长
        time_step = self.model_info.get('time_step', self.time_step) if self.model_info else self.time_step

        # 准备预测数据
        scaled_data = self.scaler.transform(self.df[['pressure']])
        last_sequence = scaled_data[-time_step:]
        last_sequence = last_sequence.reshape(1, time_step, 1)

        # 预测未来几天
        future_dates = []
        future_predictions = []

        current_sequence = last_sequence.copy()

        print(f"开始预测未来 {days} 天气压情况...")
        for i in range(days):
            # 预测下一天气压
            next_pressure = self.model.predict(current_sequence, verbose=0)[0][0]
            future_predictions.append(next_pressure)

            # 更新序列用于下一次预测
            temp_seq = current_sequence[0].copy()
            temp_seq = np.roll(temp_seq, -1, axis=0)
            temp_seq[-1] = next_pressure
            current_sequence = temp_seq.reshape(1, time_step, 1)

            # 生成下一天日期
            next_date = self.df.index[-1] + pd.Timedelta(days=i + 1)
            future_dates.append(next_date)

        # 反向转换预测结果
        future_predictions_array = np.array(future_predictions).reshape(-1, 1)
        future_pressure = self.scaler.inverse_transform(future_predictions_array)

        # 创建未来预测数据框
        future_df = pd.DataFrame({
            'date': future_dates,
            'predicted_pressure': future_pressure.flatten()
        })
        future_df.set_index('date', inplace=True)

        # 保存预测结果
        csv_path = f'{self.base_filename}.csv'
        future_df.to_csv(csv_path)
        print(f"预测结果已保存到 {csv_path}")

        return future_df

    def visualize_results(self, future_df, history_days=30):
        """可视化预测结果"""
        if future_df is None or self.df is None:
            print("没有可视化的数据")
            return

        days = len(future_df)

        # 创建图表
        fig, ax = plt.subplots(figsize=(14, 8))

        # 历史气压
        plt.plot(self.df.index[-history_days:], self.df['pressure'][-history_days:],
                 'o-', color='royalblue', alpha=0.7, linewidth=2, markersize=5,
                 label='历史气压')

        # 预测气压
        plt.plot(future_df.index, future_df['predicted_pressure'],
                 'o-', color='crimson', alpha=0.8, linewidth=2.5, markersize=8,
                 label='预测气压')

        # 添加数据标签
        for i, (date, pressure) in enumerate(zip(future_df.index, future_df['predicted_pressure'])):
            plt.text(date, pressure + 1, f'{pressure:.1f}',
                     ha='center', va='bottom', color='darkred', fontsize=10, fontweight='bold')

        # 设置图表格式
        plt.title('常州市未来气压预测', fontsize=18, fontweight='bold', pad=15)
        plt.xlabel('日期', fontsize=14)
        plt.ylabel('气压 (hPa)', fontsize=14)
        plt.grid(True, linestyle='--', alpha=0.7)

        # 添加标准大气压参考线
        ax.axhline(y=1013.25, color='forestgreen', linestyle='--', alpha=0.7)
        plt.text(self.df.index[-history_days], 1013.25 + 0.5, '标准大气压 (1013.25 hPa)',
                 color='forestgreen', fontsize=10)

        # 添加预测区域背景
        min_date = self.df.index[-1]
        ax.axvspan(min_date, future_df.index[-1], alpha=0.1, color='pink', label='预测区域')

        # 格式化日期
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%m-%d'))
        ax.xaxis.set_major_locator(mdates.DayLocator(interval=5))
        fig.autofmt_xdate()

        # 添加数据统计信息
        avg_history = self.df['pressure'][-history_days:].mean()
        avg_prediction = future_df['predicted_pressure'].mean()
        max_prediction = future_df['predicted_pressure'].max()
        min_prediction = future_df['predicted_pressure'].min()

        stat_text = f"历史平均: {avg_history:.1f} hPa\n"
        stat_text += f"预测平均: {avg_prediction:.1f} hPa\n"
        stat_text += f"预测最高: {max_prediction:.1f} hPa\n"
        stat_text += f"预测最低: {min_prediction:.1f} hPa"

        # 在图表右上角添加统计信息
        plt.annotate(stat_text, xy=(0.97, 0.97), xycoords='axes fraction',
                     fontsize=11, ha='right', va='top',
                     bbox=dict(boxstyle="round,pad=0.5", fc="lightyellow", alpha=0.8))

        # 图例
        plt.legend(loc='upper left', fontsize=12, framealpha=0.8)

        plt.tight_layout()
        plt.savefig(f'{self.base_filename}_prediction.png', dpi=300, bbox_inches='tight')
        plt.show()

    def run(self, days=7, force_train=False):
        """运行完整的预测流程"""
        # 加载数据
        self.load_data()

        # 训练或加载模型
        if force_train or not os.path.exists(self.model_path) or not self.load_model():
            self.train_model()

        # 预测未来气压
        future_df = self.predict(days=days)

        # 可视化结果
        self.visualize_results(future_df)

        return future_df


# 主函数示例
if __name__ == "__main__":
    # 创建预测器实例
    forecaster = PressureForecaster(data_file='./merged_output/5.csv')

    # 运行完整流程
    predictions = forecaster.run(days=7)

    print("\n未来气压预测结果:")
    print(predictions)